package com.techyuan.algorithm.structure.tree;

import com.techyuan.algorithm.utils.PrintUtil;
import com.techyuan.algorithm.utils.SortUtils;
import com.techyuan.algorithm.utils.test.BaseTest;
import com.techyuan.algorithm.utils.test.TestUtils;

import java.util.Map;
import java.util.TreeMap;

/**
 * TODO 红黑树
 * 红黑树特性：
 * 1、每个结点要么是红的，要么是黑的
 * 2、根节点是黑色
 * 3、每个叶子节点都是黑色的
 * 4、如果一个节点是红色的那么子节点都是黑色
 * 5、对每个结点，从该结点到其子孙结点的所有路径上包含相同数目的黑结点
 * Created by tianjuntao on 2017/1/9.
 */
public class RBTree extends BSTree {
    public RBTree() {
        super(RBNode.nil);
    }

    @Override
    public void insert(BSTree tree, Integer z) {
        RBNode insertNode = new RBNode(RBNode.nil, RBNode.nil, RBNode.nil, z, Color.RED);
        insert(tree, insertNode);
        insertFixup(tree, insertNode);
    }

    /**
     * 红黑树插入重着色和调整程序
     *
     * @param tree
     * @param z
     */
    private void insertFixup(BSTree tree, RBNode z) {
        while (z.getParent().getColor() == Color.RED) {
            if (z.getParent() == z.getParent().getParent().getLeft()) {
                RBNode y = z.getParent().getParent().getRight();
                if (y.getColor() == Color.RED) {                     //case 1
                    z.getParent().setColor(Color.BLACK);                                                //case 1
                    z.getParent().getParent().setColor(Color.RED);                                      //case 1
                    y.setColor(Color.BLACK);                         //case 1
                    z = z.getParent().getParent();                                                      //case 1
                } else {
                    if (z == z.getParent().getRight()) {
                        z = z.getParent();                                                              //case 2
                        leftRotate(tree, z);                                                //case 2
                    }
                    z.getParent().setColor(Color.BLACK);                                                //case 3
                    z.getParent().getParent().setColor(Color.RED);                                      //case 3
                    rightRotate(tree, z.getParent().getParent());                                       //case 3
                }
            } else {
                if (z.getParent().getParent().getLeft().getColor() == Color.RED) {                     //case 1
                    z.getParent().setColor(Color.BLACK);                                                //case 1
                    z.getParent().getParent().setColor(Color.RED);                                      //case 1
                    z.getParent().getParent().getLeft().setColor(Color.BLACK);                         //case 1
                    z = z.getParent().getParent();                                                      //case 1
                } else {
                    if (z == z.getParent().getLeft()) {
                        z = z.getParent();                                                              //case 2
                        rightRotate(tree, z);                                                //case 2
                    }
                    z.getParent().setColor(Color.BLACK);                                                //case 3
                    z.getParent().getParent().setColor(Color.RED);                                      //case 3
                    leftRotate(tree, z.getParent().getParent());                                       //case 3
                }
            }
        }
        ((RBNode) tree.root).setColor(Color.BLACK);
    }

    /**
     * 左旋
     *
     * @param tree
     * @param x
     */
    protected void leftRotate(BSTree tree, TreeNode x) {
        TreeNode y = x.getRight();
        //先构建x的子树
        x.setRight(y.getLeft());
        if (!x.getRight().isNil()) {
            x.getRight().setParent(x);
        }
        //y和父节点的关系
        y.setParent(x.getParent());
        if (x.getParent().isNil())
            tree.root = y;
        else if (x == x.getParent().getRight())
            x.getParent().setRight(y);
        else
            x.getParent().setLeft(y);
        //将x作为y的左孩子
        y.setLeft(x);
        x.setParent(y);
    }

    /**
     * 右旋
     *
     * @param tree
     * @param x
     */
    protected void rightRotate(BSTree tree, TreeNode x) {
        TreeNode y = x.getLeft();
        //处理x子树
        x.setLeft(y.getRight());
        if (!y.getRight().isNil()) {
            y.getRight().setParent(x);
        }
        //处理y和父节点的关系
        y.setParent(x.getParent());
        if (x.getParent().isNil())//树根
            tree.root = y;
        else if (x == x.getParent().getRight())
            x.getParent().setRight(y);
        else
            x.getParent().setLeft(y);
        //设置x为y的左孩子
        y.setRight(x);
        x.setParent(y);
    }

    @Override
    public TreeNode deleteNode(BSTree tree, TreeNode z) {
        RBNode deletedNode = (RBNode) z;
        RBNode x;
        Color originalColor = deletedNode.color;
        if (deletedNode.getLeft().isNil()) {
            x = deletedNode.getRight();
            transplant(tree, deletedNode, deletedNode.getRight());
        } else if (deletedNode.getRight().isNil()) {
            x = deletedNode.getLeft();
            transplant(tree, deletedNode, deletedNode.getLeft());
        } else {
            RBNode y = (RBNode) successor(deletedNode);
            originalColor = y.color;
            x = y.getRight();
            if (y.getParent() != deletedNode) {
                transplant(tree, y, y.getRight());
                y.setRight(((RBNode) z).getRight());
                y.getRight().setParent(y);
            }
            transplant(tree, deletedNode, y);
            y.setLeft(deletedNode.getLeft());
            y.getLeft().setParent(y);
            y.setColor(deletedNode.getColor());
        }
        if (originalColor == Color.BLACK) {
            rbDeleteFixup((RBTree) tree, x);
        }
        return deletedNode;
    }

    public void rbDeleteFixup(RBTree tree, RBNode x) {
        while (x != root && x.getColor() == Color.BLACK) {
            if (x == x.getParent().getLeft()) {
                RBNode sib = x.getParent().getRight();

                if (sib.getColor() == Color.RED) {
                    sib.setColor(Color.BLACK);
                    x.getParent().setColor(Color.RED);
                    leftRotate(tree, x.getParent());
                    sib = x.getParent().getRight();
                }

                if (x.getLeft().getColor() == Color.BLACK &&
                        x.getRight().getColor() == Color.BLACK) {
                    sib.setColor(Color.RED);
                    x = x.getParent();
                } else {
                    if (x.getRight().getColor() == Color.BLACK) {
                        sib.getLeft().setColor(Color.BLACK);
                        sib.setColor(Color.RED);
                        rightRotate(tree, sib);

                        sib = x.getParent().getRight();
                    }
                    sib.setColor(x.getParent().getColor());
                    x.getParent().setColor(Color.BLACK);
                    x.getRight().setColor(Color.BLACK);
                    leftRotate(tree, x.getParent());
                    x = (RBNode) root;
                }
            } else { // symmetric
                RBNode sib = x.getParent().getLeft();

                if (sib.getColor() == Color.RED) {
                    sib.setColor(Color.BLACK);
                    x.getParent().setColor(Color.RED);
                    rightRotate(tree, x.getParent());
                    sib = x.getParent().getLeft();
                }

                if (sib.getRight().getColor() == Color.BLACK &&
                        sib.getLeft().getColor() == Color.BLACK) {
                    sib.setColor(Color.RED);
                    x = x.getParent();
                } else {
                    if (x.getLeft().getColor() == Color.BLACK) {
                        sib.getRight().setColor(Color.BLACK);
                        sib.setColor(Color.RED);
                        leftRotate(tree, sib);
                        sib = x.getParent().getLeft();
                    }
                    sib.setColor(x.getParent().getColor());
                    x.getParent().setColor(Color.BLACK);
                    sib.getLeft().setColor(Color.BLACK);
                    rightRotate(tree, x.getParent());
                    x = (RBNode) root;
                }
            }
        }

        x.setColor(Color.BLACK);
    }

    public static class RBNode extends TreeNode<RBNode> {
        public static RBNode nil;

        static {
            nil = new RBNode(null);
            nil.setParent(nil);
            nil.setLeft(nil);
            nil.setRight(nil);
            nil.setKey(Integer.MIN_VALUE);
            nil.setColor(Color.BLACK);
        }

        public RBNode() {
            super();
            color = Color.BLACK;
        }

        public RBNode(Integer key) {
            super(key);
            color = Color.BLACK;
        }

        public RBNode(RBNode l, RBNode r, RBNode p, Integer k, Color color) {
            super(l, r, p, k);
            this.color = color;
        }

        private Color color;

        public void setColor(Color color) {
            this.color = color;
        }

        public Color getColor() {
            return color;
        }

        public boolean isNil() {
            return this.getKey().equals(Integer.MIN_VALUE);
        }

        public TreeNode getNil() {
            return nil;
        }

        @Override
        public String toString() {
            return String.format("[%d,%s]", this.getKey(), this.getColor());
        }
    }

    public enum Color {
        BLACK, RED;
    }

    public static void main(String[] args) {
        int[] arr = {41, 38, 31, 12, 19, 8}; //{41, 16, 81, 47, 21, 60, 0, 93, 34};//SortUtils.createRandomList(10, 100);
        PrintUtil.printArrPretty(arr);
        RBTree tree = new RBTree();
        for (int v : arr)
            tree.insert(tree, v);
        TreeUtils.preorderTreeWalk(tree.root);//先序
        System.out.println("树节点数:" + TreeUtils.getTreeSize(tree.root, tree.root.nil));
        int[] delArr = {8, 12, 19, 31, 38, 41};
        for (int v : delArr) {
            tree.deleteNode(tree, v);
            System.out.println();
            TreeUtils.preorderTreeWalk(tree.root);//先序
        }
        TreeMap<Integer, Integer> treeMap = new TreeMap<Integer, Integer>();
        for (int v : arr)
            treeMap.put(v, v);
        Map.Entry<Integer, Integer> treeMapRoot = treeMap.firstEntry();
        treeMap.remove(41);
        tree.deleteNode(tree, 41);
        TreeUtils.preorderTreeWalk(tree.root);//先序
        TreeUtils.preorderTreeWalk(tree.root);//先序
    }
}
